import os
import cv2
import face_recognition

# 模型数据图片目录
filePath = "../img/face_recognition"
cap = cv2.VideoCapture(0)
total_image_name = []
total_face_encoding = []

for fn in os.listdir(filePath):
    # fn表示文件名
    total_face_encoding.append(
        face_recognition.face_encodings(face_recognition.load_image_file(filePath + "/" + fn))[0]
    )

    # 截取图片名（这里应该把images文件中的图片名命名为为人物名）
    fn = fn[:(len(fn) - 4)]
    # 图片名字列表
    total_image_name.append(fn)

while 1:
    ret, frame = cap.read()

    # 发现在视频帧中所有的脸和face_encodings
    face_locations = face_recognition.face_locations(frame)
    face_encodings = face_recognition.face_encodings(frame, face_locations)
    name = "Unknown"

    # 在这个视频中循环遍历每个人脸
    for (top, right, bottom, left), face_encoding in zip(face_locations, face_encodings):
        # 看面部是否和已知人脸匹配
        for i, v in enumerate(total_face_encoding):
            match = face_recognition.compare_faces([v], face_encoding, tolerance=0.5)

            if match[0]:
                name = total_image_name[i]
                break

        # 画出一个框，框住脸
        cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), 2)
        # 画出一个带名字标签，放框下
        cv2.rectangle(frame, (left, bottom - 35), (right, bottom), (0, 0, 255), cv2.FILLED)

        font = cv2.FONT_HERSHEY_DUPLEX
        cv2.putText(frame, name, (left + 6, bottom - 6), font, 1.0, (255, 255, 255), 1)

    # 显示结果图像
    cv2.imshow('Video', frame)
    if cv2.waitKey(1) & 0xff == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()
